package com.zjj.learn.algorithmPlus.segment2.binaryTree;

/**
 * 平衡二插搜索树
 *
 * @author zjj_admin
 */
public class AVLTree<K extends Comparable<K>, V> {

    /**
     * 根节点
     */
    AVLNode<K, V> root;

    public AVLTree() {
    }


    public void preOrder() {
        preOrder(root);
    }

    public void inOrder() {
        inOrder(root);
    }

    public void postOrder() {
        postOrder(root);
    }

    /**
     * 前序遍历，继续递归实现
     *
     * @param root
     */
    private void preOrder(AVLNode<K, V> root) {
        if (root == null) {
            return;
        }
        System.out.print(root.key + "" + root.data + "\t");
        preOrder(root.left);
        preOrder(root.right);
    }

    private void inOrder(AVLNode<K, V> root) {
        if (root == null) {
            return;
        }
        inOrder(root.left);
        System.out.print(root.key + "" + root.data + "\t");
        inOrder(root.right);
    }

    private void postOrder(AVLNode<K, V> root) {
        if (root == null) {
            return;
        }
        postOrder(root.left);
        postOrder(root.right);
        System.out.print(root.key + "" + root.data + "\t");
    }

    /**
     * 添加一个节点
     * 使用递归的思想添加节点
     *
     * @param node
     */
    public void put(AVLNode<K, V> node) {
        root = doPut(root, node);
    }

    /**
     * 递归添加操作
     *
     * @param parent 父节点
     * @param node   添加节点
     * @return
     */
    private AVLNode<K, V> doPut(AVLNode<K, V> parent, AVLNode<K, V> node) {
        if (parent == null) {
            return node;
        }
        int compare = node.key.compareTo(parent.key);
        if (compare == 0) {
            // 修改数据
            parent.data = node.data;
        } else if (compare < 0) {
            parent.left = doPut(parent.left, node);
        } else {
            parent.right = doPut(parent.right, node);
        }
        //修改父节点的高度
        updateHeight(parent);
        //将父节点进行平衡操作并直接返回
        return balance(parent);
    }

    /**
     * 删除指定 key 的节点
     *
     * @param key
     */
    @SuppressWarnings("all")
    public void remove(K key) {
        AVLNode<K, V> curr = root;
        // 父节点
        AVLNode<K, V> parent = null;
        while (curr != null) {
            int compare = key.compareTo(curr.key);
            if (compare < 0) {
                parent = curr;
                curr = curr.left;
            } else if (compare > 0) {
                parent = curr;
                curr = curr.right;
            } else {
                break;
            }
        }
        //当指定的 key 不存在时，直接返回即可
        if (curr != null) {
            //删除节点
            remove(parent, curr);
            parent = balance(parent);
            updateHeight(parent);
        }


    }

    /**
     * 删除节点
     *
     * @param parent  待删除节点的父节点
     * @param removed 待删除的节点
     */
    @SuppressWarnings("all")
    private void remove(AVLNode<K, V> parent, AVLNode<K, V> removed) {
        // 1、当被删除节点的两个子节点都为 null 时
        if (removed.left == null && removed.right == null) {
            if (parent.left == removed) {
                parent.left = null;
            } else {
                //当被删除节点为父节点的右节点时
                parent.right = null;
            }
            return;
        }
        // 2、当被删除节点的左节点为 null，右节点不为 null 时
        if (removed.left == null && removed.right != null) {
            //当被删除节点为父节点的左节点时
            if (parent.left == removed) {
                parent.left = removed.right;
                removed.right = null;
            } else {
                //当被删除节点为父节点的右节点时
                parent.right = removed.right;
                removed.right = null;
            }
            return;
        }
        // 3、当被删除节点的右节点为 null ，左节点不为 null 时
        if (removed.left != null && removed.right == null) {
            //当被删除节点为父节点的左节点时
            if (parent.left == removed) {
                parent.left = removed.left;
                removed.right = null;
            } else {
                //当被删除节点为父节点的右节点时
                parent.right = removed.left;
                removed.right = null;
            }
            return;
        }
        //4、当被删除节点的左节点和右节点都不为 null 时
        if (removed.left != null && removed.right != null) {
            //找到被删除节点的后继节点
            AVLNode<K, V> houji = removed.right;
            //后继节点的父节点
            AVLNode<K, V> houjiParent = houji;
            while (houji.left != null) {
                houjiParent = houji;
                houji = houji.left;
            }
            //当后继节点和被删除节点相邻时
            if (houjiParent == houji) {
                // 使用后继节点替换被删除节点即可
                removed.data = houji.data;
                removed.key = houji.key;
                //后继节点不可能有左子节点，这是由线索二叉树的性质决定的
                removed.right = houji.right;
            } else {
                //被删除节点的值就是后继节点
                removed.data = houji.data;
                removed.key = houji.key;
                //删除后继节点
                remove(houjiParent, houji);
                parent = balance(parent);
                updateHeight(houjiParent);
            }
        }
    }


    /**
     * 求一个节点的高度
     *
     * @param node
     * @return
     */
    private int height(AVLNode<K, V> node) {
        return node == null ? 0 : node.height;
    }

    /**
     * 更新一个节点的高度
     *
     * @param node
     * @return
     */
    private void updateHeight(AVLNode<K, V> node) {
        node.height = Math.max(height(node.left), height(node.right)) + 1;
    }

    /**
     * 平衡因子 = 左子树高度 - 右子树高度
     *
     * @param node
     * @return
     */
    private int balanceFactor(AVLNode<K, V> node) {
        return height(node.left) - height(node.right);
    }

    /**
     * 左旋操作
     * 需要左旋的前提：
     * 1、节点的左子树高度 - 右子树高度 > 1
     * 2、对于节点的左孩子节点。左子树高度 >= 右子树高度
     *
     * @param node 需要旋转的节点
     * @return
     */
    private AVLNode<K, V> leftRoute(AVLNode<K, V> node) {
        if (node.left == null) {
            AVLNode<K, V> newRoot = node.right;
            node.right = null;
            newRoot.left = node;
            return newRoot;
        }
        //节点左节点成为新的根节点
        AVLNode<K, V> newRoot = node.left;
        //让左节点的右孩子指向根节点
        node.left = newRoot.right;
        //新根节点的右孩子指向原始的根节点
        newRoot.right = node;
        //更新根节点和新的根节点的高度的高度
        //两个更新操作不能交换顺序
        updateHeight(node);
        updateHeight(newRoot);
        if(node == root){
            root = newRoot;
        }
        //返回新的根节点
        return newRoot;
    }

    /**
     * 右旋操作
     * 需要左旋的前提：
     * 1、节点的右子树高度 - 左子树高度 > 1
     * 2、对于节点的右孩子节点。左子树高度 <= 右子树高度
     *
     * @param node 需要旋转的节点
     * @return
     */
    private AVLNode<K, V> rightRoute(AVLNode<K, V> node) {
        if (node.right == null) {
            AVLNode<K, V> newRoot = node.left;
            node.left = null;
            newRoot.right = node;
            return newRoot;
        }
        //节点右节点成为新的根节点
        AVLNode<K, V> newRoot = node.right;
        //让新的根节点的左孩子成为根节点的右孩子
        node.right = newRoot.left;
        //新根节点的左孩子指向原始的根节点
        newRoot.left = node;
        //更新根节点和新的根节点的高度的高度
        //两个更新操作不能交换顺序
        updateHeight(node);
        updateHeight(newRoot);
        if(node == root){
            root = newRoot;
        }
        //返回新的根节点
        return newRoot;
    }


    /**
     * 左右旋操作
     * 需要左右旋的前提：
     * 1、节点的左子树高度 - 右子树高度 > 1
     * 2、对于节点的左孩子节点。左子树高度 < 右子树高度
     * <p>
     * 旋转步骤：
     * 先让左孩子节点左旋，然后再让根节点右旋
     *
     * @param node 需要旋转的节点
     * @return
     */
    private AVLNode<K, V> leftRightRoute(AVLNode<K, V> node) {
        AVLNode<K, V> newLeft = leftRoute(node.left);
        node.left = newLeft;
        AVLNode<K, V> newRoot = rightRoute(node);
        return newRoot;
    }


    /**
     * 右左旋操作
     * 需要右左旋转的前提：
     * 1、节点的右子树高度 - 左子树高度 > 1
     * 2、对于节点的右孩子节点。左子树高度 > 右子树高度
     * <p>
     * 旋转步骤：
     * 先让节点的右孩子节点右旋，在让节点左旋
     *
     * @param node
     * @return
     */
    private AVLNode<K, V> rightLeftRoute(AVLNode<K, V> node) {
        AVLNode<K, V> newRight = rightRoute(node.right);
        node.right = newRight;
        AVLNode<K, V> newRoot = leftRoute(node);
        return newRoot;
    }


    /**
     * 平衡操作，当 node 不平衡时，让其平衡
     *
     * @param node
     * @return
     */
    private AVLNode<K, V> balance(AVLNode<K, V> node) {
        int factor = balanceFactor(node);
        AVLNode<K, V> newRoot = node;
        if (factor > 1 && balanceFactor(newRoot.left) >= 0) {
            // L
            newRoot = rightRoute(newRoot);
        } else if (factor > 1 && balanceFactor(newRoot.left) < 0) {
            // LR
            newRoot = leftRightRoute(newRoot);
        } else if (factor < -1 && balanceFactor(node.right) <= 0) {
            // R
            newRoot = rightRoute(newRoot);
        } else if (factor < -1 && balanceFactor(node.right) > 0) {
            // RL
            newRoot = rightLeftRoute(newRoot);
        }
        return newRoot;
    }


}
